import networkx as nx
import pandas as pd
import os
import matplotlib.pyplot as plt
from matplotlib import font_manager
from typing import Dict, List, Tuple

class StockGraphConstructor:
    def __init__(self, data_dir: str = "hs300_data"):
        """
        初始化图模型构建器
        :param data_dir: 预处理数据存放目录（包含graph_*.csv文件）
        """
        self.data_dir = data_dir
        self.G = nx.Graph()  # 异构图（股票节点+行业节点）
        self.industry_corr: Dict[Tuple[str, str], float] = {}  # 行业间相关度
        
        # 检查数据文件是否存在
        required_files = [
            "graph_stock_nodes.csv",
            "graph_stock_links.csv",
            "graph_industry_links.csv"
        ]
        for file in required_files:
            if not os.path.exists(os.path.join(data_dir, file)):
                raise FileNotFoundError(f"缺少必要数据文件：{os.path.join(data_dir, file)}")
        print(f"数据文件检查通过，开始构建图模型...")

    def load_and_construct_graph(self) -> nx.Graph:
        """加载数据并构建异构图（新增股票名称属性）"""
        # 1. 加载股票节点数据（含属性：波动率、平均成交额、行业、股票名称）
        stock_nodes = pd.read_csv(os.path.join(self.data_dir, "graph_stock_nodes.csv"), encoding="utf-8-sig")
        # 验证股票名称列是否存在
        if "stock_name" not in stock_nodes.columns:
            raise ValueError("graph_stock_nodes.csv 缺少'stock_name'列，请添加股票名称列或修改代码中列名")
        print(f"加载股票节点数据：共{len(stock_nodes)}只股票（已包含股票名称属性）")
        
        # 添加股票节点到图中（新增stock_name属性）
        for _, row in stock_nodes.iterrows():
            self.G.add_node(
                row["stock_code"],  # 节点ID：股票代码
                type="stock",       # 节点类型：stock（股票）
                volatility=row["volatility"],  # 波动率（风险指标）
                avg_turnover=row["avg_turnover_billion"],  # 平均成交额（亿元）
                industry=row["industry"],  # 所属行业
                stock_name=row["stock_name"]  # 股票名称（用于可视化标签）
            )
        
        # 2. 加载行业节点（从股票-行业归属边中提取）
        industry_links = pd.read_csv(os.path.join(self.data_dir, "graph_industry_links.csv"), encoding="utf-8-sig")
        industries = industry_links["industry"].unique()
        print(f"加载行业节点数据：共{len(industries)}个行业")
        print("所有行业节点：", industries)
        
        # 添加行业节点到图中
        for industry in industries:
            self.G.add_node(
                industry,
                type="industry"  # 节点类型：industry（行业）
            )
        
        # 3. 加载股票-股票联动边（带相关系数权重）
        stock_links = pd.read_csv(os.path.join(self.data_dir, "graph_stock_links.csv"), encoding="utf-8-sig")
        print(f"加载股票联动边数据：共{len(stock_links)}条边")
        
        # 添加股票联动边到图中
        for _, row in stock_links.iterrows():
            self.G.add_edge(
                row["stock1"],
                row["stock2"],
                type="stock_corr",  # 边类型：股票价格相关
                weight=row["corr"]  # 权重：相关系数（-1到1）
            )
        
        # 4. 加载股票-行业归属边（权重固定为1）
        print(f"加载股票-行业归属边数据：共{len(industry_links)}条边")
        for _, row in industry_links.iterrows():
            self.G.add_edge(
                row["stock_code"],
                row["industry"],
                type="stock_industry",  # 边类型：股票属于某行业
                weight=1.0  # 归属关系权重固定为1
            )
        
        return self.G

    def validate_graph_integrity(self) -> None:
        """验证图结构完整性，输出关键统计信息"""
        # 统计节点数量
        total_nodes = self.G.number_of_nodes()
        stock_nodes = [n for n, d in self.G.nodes(data=True) if d["type"] == "stock"]
        industry_nodes = [n for n, d in self.G.nodes(data=True) if d["type"] == "industry"]
        
        # 检查股票节点是否包含名称属性
        stock_without_name = [n for n, d in self.G.nodes(data=True) if d["type"] == "stock" and "stock_name" not in d]
        if stock_without_name:
            print(f"警告：{len(stock_without_name)}个股票节点缺少名称属性（可能数据异常）")
        else:
            print(f"所有{len(stock_nodes)}个股票节点均包含名称属性，可正常显示代码+名称")
        
        # 统计边数量
        total_edges = self.G.number_of_edges()
        stock_corr_edges = [e for e in self.G.edges(data=True) if e[2]["type"] == "stock_corr"]
        stock_industry_edges = [e for e in self.G.edges(data=True) if e[2]["type"] == "stock_industry"]
        
        # 输出验证结果
        print("\n=== 图模型完整性验证 ===")
        print(f"总节点数：{total_nodes}（股票节点：{len(stock_nodes)}，行业节点：{len(industry_nodes)}）")
        print(f"总边数：{total_edges}（股票联动边：{len(stock_corr_edges)}，股票-行业归属边：{len(stock_industry_edges)}）")
        print(f"平均每个股票节点的关联边数：{len(stock_corr_edges)/len(stock_nodes):.2f}")
        print(f"平均每个行业包含的股票数：{len(stock_industry_edges)/len(industry_nodes):.2f}")
        
        # 检查是否有孤立节点
        isolated_nodes = list(nx.isolates(self.G))
        if isolated_nodes:
            print(f"警告：存在{len(isolated_nodes)}个孤立节点（无任何边连接），可能是数据清洗不彻底导致")
        else:
            print("图模型无孤立节点，结构完整")

    def calculate_industry_correlation(self) -> None:
        """计算行业间相关度"""
        # 1. 构建行业-股票映射表
        industry_stocks: Dict[str, List[str]] = {}
        for node in self.G.nodes:
            if self.G.nodes[node]["type"] == "stock":
                industry = self.G.nodes[node]["industry"]
                if industry not in industry_stocks:
                    industry_stocks[industry] = []
                industry_stocks[industry].append(node)
        
        # 2. 计算每对行业的平均相关度
        industries = list(industry_stocks.keys())
        print(f"\n开始计算行业间相关度（共{len(industries)}个行业，需计算{len(industries)*(len(industries)-1)//2}对）...")
        
        for i, ind1 in enumerate(industries):
            for j in range(i + 1, len(industries)):
                ind2 = industries[j]
                stocks1 = industry_stocks[ind1]
                stocks2 = industry_stocks[ind2]
                
                # 跳过股票数量过少的行业
                if len(stocks1) < 3 or len(stocks2) < 3:
                    continue
                
                # 计算两个行业内股票的平均联动系数
                corr_scores = []
                for s1 in stocks1[:5]:
                    for s2 in stocks2[:5]:
                        if self.G.has_edge(s1, s2):
                            corr_scores.append(self.G[s1][s2]["weight"])
                
                # 计算平均相关度并添加到图中
                if corr_scores:
                    avg_corr = sum(corr_scores) / len(corr_scores)
                    self.industry_corr[(ind1, ind2)] = avg_corr
                    self.G.add_edge(
                        ind1,
                        ind2,
                        type="industry_corr",
                        weight=avg_corr
                    )
        
        print(f"行业间相关度计算完成，共添加{len(self.industry_corr)}条行业-行业边")
        
        # 保存行业相关度数据
        industry_corr_df = pd.DataFrame([
            {"industry1": k[0], "industry2": k[1], "correlation": v}
            for k, v in self.industry_corr.items()
        ])
        industry_corr_df.to_csv(
            os.path.join(self.data_dir, "industry_correlation.csv"),
            encoding="utf-8-sig",
            index=False
        )
        print(f"行业相关度数据已保存至：{os.path.join(self.data_dir, 'industry_correlation.csv')}")

    def _get_valid_chinese_font(self) -> str:
        """检测并返回系统中可用的中文字体"""
        # 1. 定义各系统常用中文字体列表（按优先级排序）
        candidate_fonts = [
            "SimHei",          # Windows 默认中文字体
            "Microsoft YaHei", # Windows 常用中文字体
            "PingFang SC",     # macOS 默认中文字体
            "Heiti SC",        # macOS 常用中文字体
            "WenQuanYi Zen Hei",# Linux 常用中文字体
            "DejaVu Sans"      # 兜底：无中文字体时用英文无衬线字体
        ]
        
        # 2. 检测候选字体是否存在于系统中
        available_fonts = set([f.name for f in font_manager.fontManager.ttflist])
        for font in candidate_fonts:
            if font in available_fonts:
                print(f"✅ 检测到可用中文字体：{font}，将用于中文渲染")
                return font
        
        # 3. 若系统无任何中文字体：手动加载本地字体文件
        custom_font_path = os.path.join(self.data_dir, "fonts", "SimHei.ttf")
        if os.path.exists(custom_font_path):
            font_manager.fontManager.addfont(custom_font_path)
            font_manager._rebuild()  # 重建字体缓存
            print(f"✅ 系统无中文字体，已手动加载本地字体：{custom_font_path}")
            return "SimHei"
        
        # 4. 终极兜底：无任何中文字体时，用英文字体
        print("⚠️  系统无任何可用中文字体，中文将显示为英文/符号，建议手动加载SimHei.ttf")
        return "DejaVu Sans"

    def visualize_sample_graph(self, sample_industry: str = "白酒", top_n_stocks: int = 5) -> None:
        """可视化样本子图（彻底解决中文显示问题）"""
        # ---------------- 1. 获取可用中文字体并全局配置 ----------------
        chinese_font = self._get_valid_chinese_font()
        # 全局字体配置（覆盖所有matplotlib组件）
        plt.rcParams.update({
            "font.family": chinese_font,  # 全局字体
            "axes.unicode_minus": False,  # 解决负号显示异常
            "font.size": 9,               # 基础字体大小
            "axes.titlepad": 20,          # 标题与坐标轴间距
            "figure.dpi": 100             # 画布DPI（临时，保存时用300）
        })
        
        # ---------------- 2. 验证子图节点基础数据 ----------------
        # 检查行业节点是否存在
        if sample_industry not in self.G.nodes:
            print(f"错误：行业节点[{sample_industry}]不存在于图中，请使用日志中的精确行业名称")
            return
        
        # 提取行业内股票
        industry_stocks = [
            n for n, d in self.G.nodes(data=True)
            if d["type"] == "stock" and d["industry"] == sample_industry
        ][:top_n_stocks]
        if not industry_stocks:
            print(f"警告：未从行业[{sample_industry}]中提取到股票，可能行业名称不匹配")
            return
        print(f"成功提取行业[{sample_industry}]的股票：{industry_stocks}")
        
        # 提取关联股票（联动边权重前3，去重）
        related_stocks = []
        for s in industry_stocks:
            neighbors = sorted(
                [n for n in self.G.neighbors(s) if self.G.nodes[n]["type"] == "stock"],
                key=lambda x: self.G[s][x]["weight"],
                reverse=True
            )[:3]
            related_stocks.extend(neighbors)
        related_stocks = list(set(related_stocks) - set(industry_stocks))  # 去重
        print(f"提取行业内股票的关联股票：{related_stocks}")
        
        # ---------------- 3. 构建子图 ----------------
        subgraph_nodes = [sample_industry] + industry_stocks + related_stocks
        subG = self.G.subgraph(subgraph_nodes)
        print(f"子图构建完成：节点数={len(subG.nodes)}，边数={len(subG.edges)}")
        
        # ---------------- 4. 自定义节点标签（代码+中文名称） ----------------
        node_labels = {}
        for node in subG.nodes:
            node_attr = self.G.nodes[node]
            if node_attr["type"] == "industry":
                # 行业节点：显示完整中文行业名
                node_labels[node] = node
            else:
                # 股票节点：显示「代码\n中文名称」
                stock_code = node
                stock_name = node_attr["stock_name"]
                # 截断过长名称
                short_name = stock_name[:8] + "..." if len(stock_name) > 8 else stock_name
                node_labels[node] = f"{stock_code}\n{short_name}"
        
        # ---------------- 5. 绘图配置 ----------------
        plt.figure(figsize=(16, 12))  # 增大画布
        
        # 圆形布局
        pos = nx.circular_layout(subG)
        
        # 节点颜色与大小
        node_colors = []
        node_sizes = []
        for node in subG.nodes:
            if node == sample_industry:
                node_colors.append("#90EE90")  # 行业节点：浅绿
                node_sizes.append(3000)        # 行业节点更大
            elif node in industry_stocks:
                node_colors.append("#FFB6C1")  # 行业内股票：浅红
                node_sizes.append(2200)
            else:
                node_colors.append("#87CEEB")  # 关联股票：浅蓝
                node_sizes.append(2200)
        
        # ---------------- 6. 绘制图元素 ----------------
        # 6.1 绘制边
        nx.draw(
            subG, pos,
            edgelist=subG.edges(),
            edge_color="#D3D3D3",
            width=1.2,
            with_labels=False
        )
        
        # 6.2 绘制节点
        nx.draw(
            subG, pos,
            nodelist=subG.nodes(),
            node_color=node_colors,
            node_size=node_sizes,
            edgecolors="#333333",
            linewidths=1.2,
            with_labels=False
        )
        
        # 6.3 绘制节点标签
        nx.draw_networkx_labels(
            subG, pos,
            labels=node_labels,
            font_size=8,
            font_weight="bold",
            font_color="#000000",
            font_family=chinese_font,
            bbox=dict(boxstyle="round,pad=0.4", facecolor="white", alpha=0.8)
        )
        
        # 6.4 绘制边权重标签
        edge_labels = {
            (u, v): f"{subG[u][v]['weight']:.2f}"
            for u, v in subG.edges()
            if subG[u][v]["type"] == "stock_corr"
        }
        nx.draw_networkx_edge_labels(
            subG, pos,
            edge_labels=edge_labels,
            font_size=7,
            font_color="#666666",
            font_family=chinese_font,
            bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.7)
        )
        
        # ---------------- 7. 绘制标题与保存图片 ----------------
        # 修复标题字体设置方式，使用fontdict参数
        title_font = {
            'fontsize': 14,
            'fontweight': 'bold',
            'fontfamily': chinese_font
        }
        plt.title(
            f"{sample_industry}行业及关联股票子图\n（节点：{len(subG.nodes)}，边：{len(subG.edges)}）",
            fontdict=title_font,  # 使用fontdict设置字体
            pad=25  # 标题与画布顶部间距
        )
        
        # 调整边距
        plt.subplots_adjust(left=0.08, right=0.92, top=0.88, bottom=0.12)
        
        # 保存图片
        save_path = os.path.join(self.data_dir, "sample_graph.png")
        plt.savefig(
            save_path,
            dpi=300,
            bbox_inches="tight",
            facecolor="white"
        )
        print(f"样本子图已保存至：{save_path}")
        plt.close()


if __name__ == "__main__":
    # 初始化图构建器
    constructor = StockGraphConstructor(data_dir="hs300_data")
    
    # 1. 构建异构图（含股票名称属性）
    graph = constructor.load_and_construct_graph()
    
    # 2. 验证图结构
    constructor.validate_graph_integrity()
    
    # 3. 计算行业间相关度
    constructor.calculate_industry_correlation()
    
    # 4. 生成样本子图
    constructor.visualize_sample_graph(
        # 此处以C27医药制造业为例，后续可生成其他样本子图
		sample_industry="C27医药制造业",
        top_n_stocks=5
    )
    
    print("\n=== 阶段1：图模型构建与验证完成 ===")
    print("输出文件：")
    print(f"1. 行业相关度数据：{os.path.abspath(os.path.join('hs300_data', 'industry_correlation.csv'))}")
    print(f"2. 样本子图可视化：{os.path.abspath(os.path.join('hs300_data', 'sample_graph.png'))}")
